import akshare as ak
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt


# 获取市场汇总信息
def stock_sse_summary_df():
    return ak.stock_sse_summary()


# [](https://zhuanlan.zhihu.com/p/20339518843)
# 获取股票数据
def fetch_stock_data(ticker, start, end):
    stock_data = ak.stock_zh_a_hist(
        symbol=ticker, period="daily", start_date=start, end_date=end, adjust="qfq"
    )
    stock_data.rename(
        columns={
            "日期": "Date",
            "股票代码": "Code",
            "开盘": "Open",
            "收盘": "Close",
            "最高": "High",
            "最低": "Low",
            "成交量": "Volume",
            "成交额": "Turnover",
        },
        inplace=True,
    )
    stock_data["Date"] = pd.to_datetime(stock_data["Date"])
    stock_data.set_index("Date", inplace=True)
    return stock_data


# 计算动量
def calculate_momentum(data, lookback):
    data["momentum"] = data["Close"].pct_change(lookback)
    return data


# 动量策略回测
def backtest(data, lookback, initial_capital=100000):
    data = calculate_momentum(data, lookback)
    data.dropna(inplace=True)

    capital = initial_capital
    position = 0
    total_values = []
    returns = []

    for i in range(len(data)):
        if data["momentum"].iloc[i] > 0:  # 买入信号
            if capital > 0:
                position = capital / data["Close"].iloc[i]
                capital = 0
        elif position > 0:  # 卖出信号
            capital = position * data["Close"].iloc[i]
            position = 0

        total_value = capital + position * data["Close"].iloc[i]
        total_values.append(total_value)
        returns.append((total_values[-1] / total_values[-2]) - 1 if i > 0 else 0)

    data["total_value"] = total_values
    data["strategy_return"] = returns
    return data


# [](https://zhuanlan.zhihu.com/p/18524149183
# 计算趋势数据
def cal_trend(stock_data, lifecycle=20):
    stock_data["High Trend"] = stock_data["High"].rolling(window=lifecycle).max()
    stock_data["Low Trend"] = stock_data["Low"].rolling(window=lifecycle).min()

    stock_data.dropna(subset=["High Trend", "Low Trend"], inplace=True)
    return stock_data


# 有问题
def plot_trend_with_lifecycle(stock_data):

    buy_signals = stock_data.dropna(subset=["Buy Price"])
    sell_signals = stock_data.dropna(subset=["Sell Price"])
    # 绘制价格与趋势线
    plt.figure(figsize=(14, 7))
    plt.scatter(
        buy_signals.index,
        buy_signals["Buy Price"],
        label="Buy Signal",
        marker="^",
        color="green",
        alpha=0.8,
    )
    plt.scatter(
        sell_signals.index,
        sell_signals["Sell Price"],
        label="Sell Signal",
        marker="v",
        color="red",
        alpha=0.8,
    )

    plt.plot(stock_data["Close"], label="Close Price", color="blue", linewidth=1.5)
    plt.plot(
        stock_data.index,
        stock_data["High Trend"],
        label=f"High Trend Line (Lifecycle={lifecycle})",
        linestyle="--",
        color="green",
        linewidth=1,
    )
    plt.plot(
        stock_data.index,
        stock_data["Low Trend"],
        label=f"Low Trend Line (Lifecycle={lifecycle})",
        linestyle="--",
        color="red",
        linewidth=1,
    )
    plt.title(f"Trend Lines with Lifecycle: {ticker}")
    plt.xlabel("Date")
    plt.ylabel("Price")
    plt.legend()
    plt.grid(alpha=0.3)
    plt.show()
